package com.itheima.stargraph.service.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.json.JSONUtil;
import com.itheima.stargraph.dto.request.Text2ImageReqDto;
import com.itheima.stargraph.dto.respone.Text2ImageResDto;
import com.itheima.stargraph.pojo.ComfyuiModel;
import com.itheima.stargraph.pojo.ComfyuiRequestDto;
import com.itheima.stargraph.pojo.ComfyuiTask;
import com.itheima.stargraph.service.FreemarkerService;
import com.itheima.stargraph.service.OllamaService;
import com.itheima.stargraph.service.Text2ImageService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

/**
 * TODO
 *
 * @Author mr.wu
 * @Data #{DATE} #{TIME}
 */
@Service
public class Text2ImageServiceImpl implements Text2ImageService {

    @Autowired
    private OllamaService ollamaService;
    @Autowired
    private FreemarkerService freemarkerService;
    private static final String WS_CLIENT_ID="star-graph";
    /**
     * 获取ComfyuiTask任务实例
     * @param dto
     * @return
     */
    @Override
    public ComfyuiTask getComfyuiTask(Text2ImageReqDto dto) {
        //1.将请求dto对象拷贝为ComfyuiModle对象
        ComfyuiModel model = BeanUtil.toBean(dto, ComfyuiModel.class);
        //2.设置ComfyuiModle对象更多属性的值
        String prompt = ollamaService.translation(dto.getPrompt());
        String reverse = ollamaService.translation(dto.getReverse());

        model.setPrompt(prompt);
        model.setReverse(reverse);
        model.setModelName(dto.modelName());
        model.setSamplerName(dto.samplerName());
        model.setScheduler(dto.scheduler());
        model.setWidth(dto.width());
        model.setHeight(dto.height());

        //3.使用Freemarker模板生成提示词
        String workflowJson = freemarkerService.renderWorkflow(model);
        //4.封装ComfyuiTask对象

        ComfyuiRequestDto comfyuiRequestDto = new ComfyuiRequestDto(WS_CLIENT_ID,  workflowJson);
        ComfyuiTask comfyuiTask = new ComfyuiTask(null,comfyuiRequestDto);
        return comfyuiTask;
    }


    @Autowired
    private StringRedisTemplate redisTemplate;

    private static final String REDIS_KEY_COUNTER="star:graph:counter";

    private static final String REDIS_KEY_QUEUE="star:graph:counter";

    private static final String REDIS_KEY_TASK="star:graph:task:%s";

    /**
     * 将封装好的任务对象添加到队列中
     * @param comfyuiTask
     * @return
     */
    public ComfyuiTask addQueueTask(ComfyuiTask comfyuiTask){

        //1.创建Redis计数器
        Long increment = redisTemplate.opsForValue().increment(REDIS_KEY_COUNTER);

        //2.创建并添加任务到zset集合中
        redisTemplate.opsForZSet().add(REDIS_KEY_QUEUE,comfyuiTask.getId(), increment);

        //3.将添加任务的Json数据到redis的string中
        redisTemplate.opsForValue().set(String.format(REDIS_KEY_TASK,comfyuiTask.getId()), JSONUtil.toJsonStr(comfyuiTask));

        comfyuiTask.setIndex(increment);
        return comfyuiTask;
    }

    @Override
    public Text2ImageResDto prompt(Text2ImageReqDto dto) {
        //1.根据获取ComfyuiTask任务实例
        ComfyuiTask comfyuiTask = this.getComfyuiTask(dto);
        //2.添加任务到redis队列中排队
        ComfyuiTask task = this.addQueueTask(comfyuiTask);
        //3.封装响应实例
        Text2ImageResDto text2ImageResDto = new Text2ImageResDto();
        text2ImageResDto.setPid(task.getId());
        text2ImageResDto.setQueueIndex(task.getIndex());
        return text2ImageResDto;
    }
}
